import numpy as np
import copy
import torch
from torch import nn

def IoU(box1,box2):
    """
    计算交并比
    """
    width1=box1[2]-box1[0]
    width2=box2[2]-box2[0]

    height1=box1[3]-box1[1]
    height2=box2[3]-box2[1]
    # flag1=np.abs(box2[2]-box1[0])<(width2+width1)
    # flag2=np.abs(box2[3]-box1[1])<(height1+height2)

    # if flag1 and flag2:
    xa=max(box1[0],box2[0])
    ya=max(box1[1],box2[1])
    xb=min(box1[2],box2[2])
    yb=min(box1[3],box2[3])
    
    n=max(0,yb-ya)*max(0,xb-xa)

    return n/(height1*width1+width2*height2-n)
    # else:
    #     return 0.0

def shift_value(box,tar):
    core_box=[(box[2]+box[0])/2,(box[3]+box[1])/2]
    core_tar=[(tar[2]+tar[0])/2,(tar[3]+tar[1])/2]
    box_=[box[2]-box[0],box[3]-box[1]]
    tar_=[tar[2]-tar[0],tar[3]-tar[1]]

    return (core_box[0]-core_tar[0])/box_[0],(core_box[1]-core_box[1])/box_[1],np.log(tar_[0]/box_[0]),np.log(tar_[1]/box_[1])

def value_shift(pre,box):
    core_box=[(box[2]+box[0])/2,(box[3]+box[1])/2]
    box_=[box[2]-box[0],box[3]-box[1]]
    res=[pre[0]*box_[0]+core_box[0],pre[1]*box_[1]+core_box[1],np.exp(pre[2])*box_[0],np.exp(pre[3])*box_[1]]
    return res


def get_label(boxes,tars):
    boxes_shape=boxes.shape
    len_tars=len(tars)
    res=[]

    for i in range(boxes_shape[0]):
        item=[]
        for j in range(boxes_shape[1]):
            one=[]
            for p in range(len_tars):
                iou=IoU(boxes[i,j],tars[p])
                # if iou>0:
                #     print("iou is {}".format(iou))
                #     print("+"*25)
                one.append(iou)
            item.append(one)
        res.append(item)

    return np.array(res)



class Anchor:
    def __init__(self,scales=2**np.arange(3,6),ratios=[0.5,1,2]) -> None:
        self.ratios=ratios
        self.scales=scales
        self.feature_map_size=None
    

    def anchor(self):
        """
        [x-width/2,x+width/2,y-height/2,y+height/2] 即为原坐标\n
        其中width=box[2]-box[0],height=box[3]-box[1]
        """
        boxes=[]

        ## 生成初始框
        box=np.array([0,0,self.scales[0],self.scales[0]])-1

        ## 根据ratio生成框

        width=box[2]-box[0]
        

        for ratio in self.ratios:
            one=list(box)
            one[2]=width*ratio
            boxes.append(one)

        ## 根据scales生成框
        boxes=np.array(boxes)
        for scale in self.scales[1:]:
            for i in range(len(self.ratios)):
                one=(scale/self.scales[0])*boxes[i]
                boxes=np.vstack((boxes,one))
        

        return boxes

    def get_source(self,x,y,boxes,floor_rate=16):
        """
        3*64*32*54
        [xmin,ymin,xmax,ymax]
        """
        width=boxes[:,2]-boxes[:,0]
        height=boxes[:,3]-boxes[:,1]
        x=x*floor_rate
        y=y*floor_rate

        boxes[:,0]=x-width/2
        boxes[:,1]=y-height/2
        boxes[:,2]=x+width/2
        boxes[:,3]=y+height/2

        return boxes


    def __call__(self,feature_map_size,floor_rate=16):
        self.feature_map_size=feature_map_size
        batch=self.feature_map_size[0]
        fw=self.feature_map_size[2]
        fh=self.feature_map_size[3]

        boxes=self.anchor()
        res=[]

        for _ in range(batch):
            item=[]
            for i in range(fw):
                for j in range(fh):
                    item.append(self.get_source(i,j,copy.deepcopy(boxes),floor_rate))

            res.append(item)

        return np.array(res,dtype=np.float64)
                    
class FPN(nn.Module):
    def __init__(self) -> None:
        super().__init__()


if __name__=="__main__":
    Anchor_generater=Anchor(scales=2*np.arange(4,7))
    feature_map=[1,3,37,50]
    anchors=Anchor_generater(feature_map)[0]

    tars=[[20,10,100,160],[50,70,190,240]]

    labels=get_label(anchors,tars)

    

  
    